#pragma once
#include<iostream>
#include<assert.h>
using namespace std;

namespace x
{
	template<class K,class V>
	struct TreeNode
	{
		pair<K, V> _kv;
		TreeNode<K,V>* _left;
		TreeNode<K,V>* _right;
		TreeNode<K,V>* _parent;
		int _bf;

		TreeNode(const pair<K, V>& kv)
			:_kv(kv)
			, _left(nullptr)
			, _right(nullptr)
			, _parent(nullptr)
			, _bf(0)
		{}
	};

	template<class K,class V>
	class AVLTree
	{
		typedef TreeNode<K, V> Node;
		bool Insert(const pair<K, V>& kv)
		{
			if (_root == nullptr)
			{
				_root = new Node(kv);
				return true;
			}
			Node* parent = nullptr;
			Node* cur = _root;
			while (cur)
			{
				if (cur->_kv.first < kv.first)
				{
					parent = cur;
					cur = cur->_right;
				}
				else
				{
					parent = cur;
					cur = cur->_left;
				}
			}
			cur = new Node(kv);
			if (parent->_kv.first < kv.first)
			{
				parent->_right = cur;
			}
			else
			{
				parent->_left = cur;
			}
			cur->_parent = parent;

			while (parent)
			{
				if (parent->_left == cur)
					parent->_bf--;
				else
					parent->_bf++;
				if (parent->_bf == 0) break;
				else if (parent->_bf == 1 || parent->_bf == -1)
				{
					cur = parent;
					parent = parent->_parent;
				}
				else if (parent->_bf == 2 || parent->_bf == -2)
				{
					if (parent->_bf == -2 && cur->_bf == -1)
						RotateR(parent);
					else if (parent->_bf == 2 && cur->_bf == 1)
						RotateL(parent);
					else if (parent->_bf == -2 && cur->_bf == 1)
						RotateLR(parent);
					else if (parent->_bf == 2 && cur->_bf == -1)
						RotateLR(parent);
					else
						assert(false);
					break;
				}
				else
					assert(false);
			}
			return true;
		}

		void RotateR(Node* parent)
		{
			Node* subL = parent->_left;
			Node* subLR = subL->_right;
			Node* pParent = parent->_parent;
			parent->_left = subLR;
			subL->_right = parent;
			if (subLR)
				subLR->_parent = parent;
			parent->_parent = subL;
			if (parent == _root)
			{
				_root = subL;
				subL->_parent = nullptr;
			}
			else
			{
				if (pParent->_left == parent)
					pParent->_left = subL;
				else
					pParent->_right = subL;
				subL->_parent = pParent;
			}
			parent->_bf = subL->_bf = 0;
		}

		void RotateL(Node* parent)
		{
			Node* subR = parent->_right;
			Node* subRL = subR->_left;
			Node* pParent = parent->_parent;
			parent ->_parent = subR;
			subR->_left = parent;
			if (subRL)
				subRL->_parent = parent;
			parent->_parent = subR;
			if (parent == _root)
			{
				_root = subR;
				subR->_parent = nullptr;
			}
			else
			{
				if (pParent->_left == parent)
					pParent->_left = subR;
				else
					pParent->_right = subR;
				subR->_parent = pParent;
			}
			parent->_bf = subR->_bf = 0;
		}

		void RotateLR(Node* parent)
		{
			Node* subL = parent->_left;
			Node* subLR = subL->_right;
			int bf = subLR->_bf;
			RotateL(parent->_left);
			RotateR(parent);
			if (bf == 1)
			{
				subL->_bf = -1;
				subLR->_bf = 0;
				parent->_bf = 0;
			}
			else if (bf == -1)
			{
				subL->_bf = 0;
				subLR->_bf = 0;
				parent->_bf = 1;
			}
			else if (bf == 2)
			{
				subL->_bf = 0;
				subLR->_bf = 0;
				parent->_bf = 0;
			}
			else
				assert(false);
		}

		void RotateRL(Node* parent)
		{
			Node* subR = parent->_right;
			Node* subRL = subR->_left;
			int bf = subRL->_bf;
			RotateR(parent->_right);
			RotateL(parent);
			if (bf == -1)
			{
				subRL->_bf = 0;
				parent->_bf = 0;
				subR->_bf = 1;
			}
			else if (bf == 1)
			{
				subRL->_bf = 0;
				parent->_bf = -1;
				subR->_bf = 0;
			}
			else if (bf == 0)
			{
				subRL->_bf = 0;
				parent->_bf = 0;
				subR->_bf = 0;
			}
			else
				assert(false);
		}

	private:
		Node* _root;
	};
}
